function [Uavg,Vavg,tausq_v_samples,Z_factor_samples,rmse_samples] = run_modQTL_confounder_spikeslab(Y, J, burnin, ngibbs)
%
% [Uavg,Vavg,tausq_v_samples,z_factor_samples,rmse_samples] = modQTL_confounder(Y, J, burnin, ngibbs)
%
% confounder correction by sparse matrix factorization
% using blocked Gibbs sampling
%
% Y = U V
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    [nprobe,nind] = size(Y);
    SIGSQ0 = var(Y(:));
    sigsq = SIGSQ0;

    U = 0.1 * randn(nprobe,J,'single');
    Vt = zeros(nind,J,'single');
    Z_factor = ones(nind,J,'single');
    Z_probe = ones(nprobe,J,'single');
    invtausq_u = 0.01*nind;
    invtausq_v = 0.01*nprobe*ones(J,1,'single');

    tausq_v_samples = NaN(J,ngibbs,'single');
    rmse_samples = NaN(burnin+ngibbs,1,'single');
    Z_factor_samples = NaN(nind,J,ngibbs,'single');

    Uavg = zeros(size(U),'single');
    Vavg = zeros(size(Vt'),'single');

    vsq_mean = zeros(1,J,'single');
    vsq_stoch = zeros(1,J,'single');

    prior_odds = -log(J-1);

    for iter = 1:(burnin + ngibbs),

        % ================ sample V ================
        % (Y - U*VZt' + uu*vzt')' * uu;
        % (Y' - VZt*U' + vzt*uu') * uu;
        % (Y'*uu - VZt*U'*uu + vzt*uu'*uu)

        YtU = Y'*U;
        UtU = U'*U;
        VZtUtU = bsxfun(@times, Vt, Z_factor) * UtU;


        for jj = 1:J,

            % update slab
            vt_old = Vt(:,jj);
            z_old = Z_factor(:,jj);
            denom_v = UtU(jj,jj) + invtausq_v(jj);

            vvt = (YtU(:,jj) - VZtUtU(:,jj) + z_old.*vt_old*UtU(jj,jj))/denom_v;
            vsq_mean(jj) = sum(vvt.^2);
            vvt = vvt + randn(nind,1,'single') * sqrt(sigsq/denom_v);

            % update precomputation
            delta_vz = (vvt - vt_old) .* z_old;
            VZtUtU = VZtUtU + bsxfun(@times, delta_vz, UtU(:,jj).');

            Vt(:,jj) = vvt;

            % udpate spike
            mass = YtU(:,jj).*vvt - VZtUtU(:,jj).*vvt + UtU(jj,jj).*(vvt.^2).*(z_old-0.5);
            zz = 1./(1+exp(- mass/sigsq - prior_odds));
            zz = single(rand(size(zz),'single') < zz);

            Z_factor(:,jj) = zz;
            delta_vz = vvt .* (zz - z_old);
            VZtUtU = VZtUtU + bsxfun(@times, delta_vz, UtU(:,jj).');

        end

        invtausq_v = max( 0.1/nprobe, min( nprobe*10, sigsq * nind ./ vsq_mean ) );

        r = betarnd(1+sum(Z_factor(:)),nind*J+sum(1-Z_factor(:)));
        prior_odds = log(r) - log(1-r);

        VZt = bsxfun(@times, Vt, Z_factor);

        % ================ sample U ================
        % (Y - U*VZt' + u_j*v_j')*v_j
        % Y*v_j - U*VZt'*v_j + u_j*v_j'*v_j
        % Y*VZt_j - U*(V*VZt)_j + (V*VZt)_jj * u_j


        YVZt = Y*VZt;
        VVZt = VZt'*VZt;

        for jj=1:J
            tt = YVZt(:,jj);
            ss = VVZt(:,jj);

            zz = Z_probe(:,jj);

            % update slab : w/o (zz.*zz) due to different variational
            denom_u = ss(jj) + invtausq_u;

            if denom_u <= 1e-4, % avoid NaN
                denom_u = 1e-4;
            end

            uu = (tt - U*ss + zz.*U(:,jj) * ss(jj))./ denom_u;
            uu = uu + randn(nprobe,1,'single').*sqrt(sigsq./denom_u);

            U(:,jj) = uu;

            % update Z
            mass = -0.5 * uu.^2 * ss(jj) / sigsq;
            mass = mass + uu.*(tt - U*ss + (zz.*uu)*ss(jj)) /sigsq;

            zz = 1./(1+exp(-mass));
            zz = single(rand(nprobe,1,'single') < zz);

            Z_probe(:,jj) = zz;
        end

        UZ = U.*Z_probe;

        R = Y - UZ*VZt';
        rss = sum(R(:).^2);
        nobs = nprobe*nind;

        rmse = sqrt(rss / nobs);
        rmse_samples(iter) = rmse;

        sigsq = (rss + 0.1*SIGSQ0*(J*nprobe + J*nind));
        sigsq = sigsq / (nobs + J*nprobe + J*nind);

        if iter > (burnin),
            tausq_v_samples(:,iter - burnin) = 1./invtausq_v;
            Uavg = Uavg + UZ / ngibbs;
            Vavg = Vavg + VZt' / ngibbs;
            Z_factor_samples(:,:,iter - burnin) = Z_factor;
        end

        fprintf(2,'Iter = %05d, rmse = %.4f, sigma = %.4f, probe activity = %.4f, rank activity = %.4f\n', ...
                iter, rmse, sqrt(sigsq), mean(Z_probe(:)), mean(Z_factor(:)));

    end

end % end of function
